import torch
import sys
import cv2
import os
import numpy as np
from flask import Flask, request, jsonify
from pathlib import PosixPath
import pathlib
import ddddocr
import json
from datetime import datetime
import platform

from config import Config# 确保目录存在

# 修复在Linux读取Windows保存的路径
if platform.system().lower() == 'linux':
    try:
        pathlib.WindowsPath = PosixPath
    except Exception:
        pass


# 初始化 Flask 应用
app = Flask(__name__)
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

# 加载 YOLOv5 模型
model = torch.hub.load('yolov5', 'custom', path=Config.MODEL_PATH,force_reload=False ,source='local')
model.conf = Config.MODEL_CONF
ocr = ddddocr.DdddOcr()

# 类别名称（和训练时一致）
CLASS_NAMES = Config.CLASS_NAMES


@app.route('/recognize', methods=['POST'])
def recognize():
    if 'image' not in request.files:
        return jsonify({'error': 'No image uploaded'}), 400

    file = request.files['image']

    # 读取图片数据到内存
    img_bytes = file.read()
    img_array = np.frombuffer(img_bytes, np.uint8)
    img = cv2.imdecode(img_array, cv2.IMREAD_COLOR)

    # 进行检测
    results = model(img)
    boxes = results.xyxy[0].cpu().numpy()

    # 收集识别结果
    full_text = ""

    for i, (x1, y1, x2, y2, conf, cls) in enumerate(boxes):
        class_name = CLASS_NAMES[int(cls)]
        crop = img[int(y1):int(y2), int(x1):int(x2)]
        _, buffer = cv2.imencode('.png', crop)

        # OCR识别
        if class_name == "card":  # 只有card类别进行OCR
            try:
                ocr_text = ocr.classification(buffer.tobytes())
                full_text += ocr_text
            except Exception as e:
                full_text += f" [OCR error: {str(e)}]"

    # 返回结果
    result_data = {
        'text': full_text
    }

    return jsonify(result_data)

if __name__ == '__main__':
    app.run(host=Config.FLASK_HOST, port=Config.FLASK_PORT, debug=False)
